#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @File   : trt_utils.py
# @Author : yuanwenjin
# @Mail   : xxxx@mail.com
# @Date   : 2020/07/01 08:44:55
# @Docs   : engine使用工具
'''

import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def save_engine(engine, engine_path):
    buf = engine.serialize()
    with open(engine_path, 'wb') as f:
        f.write(buf)

def load_engine(engine_path):
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        engine_data = f.read()
        engine = runtime.deserialize_cuda_engine(engine_data)
    return engine
